#include "VolPathIntegrator.cuh"

using namespace OBR;

extern "C" __device__ Spectrum __continuation_callable__volpath_integrator(SamplePRD *prd)
{
    const uint mat_base = prd->params->material_base_offset;
    const uint med_base = prd->params->medium_base_offset;
    prd->consider_medium = true;

    CallableData *sbt_record = reinterpret_cast<CallableData *>(optixGetSbtDataPointer());
    VolPathIntegratorData *settings =
        reinterpret_cast<VolPathIntegratorData *>(sbt_record->callable_payload);

    // const params
    const uint max_depth = settings->max_depth;
    const uint rr_start = settings->rr_start;
    const float rr_max = settings->rr_max;
    const float rr_thres = settings->rr_thres;
    const float eps = 1e-4;

    uint u0, u1;
    packPointer(prd, u0, u1);
    Spectrum beta = Spectrum(1.0f);
    Spectrum L = Spectrum(0.0f);
    for (size_t depth = 0; depth < max_depth; depth++)
    {

#ifdef DEBUG_SHADER
        if (prd->log)
        {
            printf(">>>>>>>>>>>>>>>>>>> depth: %d\n", depth);
        }
#endif

        float4 urand = curand_uniform4(prd->rng_long);

        // trace current incomming ray
        optixTrace(                        // trace new ray
            prd->params->handle,           // handel
            prd->inter.wi.o,               // ray origin
            prd->inter.wi.d,               // ray direction
            eps, 1e10f,                    // ray time
            prd->scene_time,               // scene time
            OptixVisibilityMask(1),        // vis mask
            OPTIX_RAY_FLAG_DISABLE_ANYHIT, // ray flags
            0,                             // SBT offset
            1,                             // SBT stride
            0,                             // miss sbt
            u0, u1);

#ifdef DEBUG_SHADER
        if (prd->log)
        {
            printf("trace result:\n");
            prd->inter.print();
            printf("wo:\n");
            prd->inter.wo.print();
        }
#endif

        // handle medium interaction if wo is in medium
        if (prd->inter.wo.med_id != -1)
        {

#ifdef DEBUG_SHADER
            if (prd->log)
            {
                printf("calling medium:\n");
            }
#endif
            // if a medium interaction happens, prd->inter is updated, and has_mi is set to true
            beta *= optixDirectCall<Spectrum, SamplePRD *, MediumCallMode>(
                prd->inter.wo.med_id + med_base, prd, MediumCallMode::TRACK_INTERACTION);
        }

        // break if beta is not valid
        if (!beta.isValid())
            break;

#ifdef DEBUG_SHADER
        if (prd->log)
        {
            printf("interaction type: %s\n",
                   prd->inter.type == InteractionType::SURFACE_INTERACTION ? "surface" : "medium");
        }
#endif

        // handle surface interaction
        if (prd->inter.type == InteractionType::SURFACE_INTERACTION)
        {
            if (!prd->hit)
            {

#ifdef DEBUG_SHADER
                if (prd->log)
                {
                    printf("ray missed\n");
                }
#endif

                /* ----------------------- case missed ---------------------- */
                L += beta * prd->bg_radiance;
                break;
            }
            else
            {
                /* -------------------- case hit surface -------------------- */
                // handle emmisive
                if (prd->inter.light_id != -1 && depth != 1)
                {

#ifdef DEBUG_SHADER
                    if (prd->log)
                    {
                        printf("hit light\n");
                    }
#endif

                    float _pdf;
                    L += beta * optixDirectCall<Spectrum, SamplePRD *, float *, int, float, float,
                                                LightCallMode>(prd->inter.light_id + 1, prd, &_pdf,
                                                               prd->inter.light_instance_id, 0, 0,
                                                               LightCallMode::CALC_LI_PDF);
                    if (prd->inter.mat_id == -1)
                        break; // break when hit light without surface
                }

                // handle surface
                if (prd->inter.mat_id == -1)
                {
                    /* --------------- case boundary of medium -------------- */
                    prd->inter.wi = prd->inter.spawnRay(prd->inter.wo.d);
                    continue;
                }
                else
                {
                    /* ----------------- case actual surface ---------------- */
                    // sample direct light
                    L += beta * uniformSampleOneLight(prd, urand.x);

                    // sample material
                    float pdf;
                    Spectrum fr = optixDirectCall<Spectrum, SamplePRD *, float *, float, float>(
                        prd->inter.mat_id + mat_base, prd, &pdf, urand.z, urand.w,
                        MaterialCallMode::SAMPLE_FR);
                    beta *= fr / pdf;
                }
            }
        }

#ifdef DEBUG_SHADER
        if (prd->log)
        {
            printf("current outgoing ray:\n");
            prd->inter.wi.print();
            printf("current beta:\n");
            beta.print();
            printf("\n");
            printf("current L:\n");
            L.print();
            printf("\n");
        }
#endif

        // rr
        float max_beta_component = beta.maxComponentValue();
        if (depth > rr_start && max_beta_component < rr_thres)
        {
            float q = clamp(fmaxf(rr_max, 1.0f - max_beta_component), 0.0f, 1.0f);
            if (curand_uniform(prd->rng) < q)
                break;
            beta /= 1 - q;
        }
        beta = beta.clamp();
    }

    return L;
}